/**
 * @author: dn-jinmin/dn-jinmin
 * @doc:
 */

package toolx

import (
	"aiworkc/internal/svc"
	"aiworkc/pkg/langchain/outputparserx"
	"context"
	"encoding/json"
	"fmt"
	"github.com/tmc/langchaingo/callbacks"
	"github.com/tmc/langchaingo/documentloaders"
	"github.com/tmc/langchaingo/embeddings"
	"github.com/tmc/langchaingo/textsplitter"
	"github.com/tmc/langchaingo/vectorstores/redisvector"
	"os"
)

type KnowledgeUpdate struct {
	svc          *svc.ServiceContext
	Callback     callbacks.Handler
	outPutParser outputparserx.Structured
	store        *redisvector.Store
}

func NewKnowledgeUpdate(svc *svc.ServiceContext) *KnowledgeUpdate {
	return &KnowledgeUpdate{
		svc:      svc,
		Callback: svc.Callbacks,
		outPutParser: outputparserx.NewStructured([]outputparserx.ResponseSchema{
			{
				Name:        "path",
				Description: "the path to file",
			}, {
				Name:        "name",
				Description: "the name to file",
			}, {
				Name:        "time",
				Description: "file update time",
			},
		}),
	}
}

func (k *KnowledgeUpdate) Name() string {
	return "knowledge_update"
}

func (k *KnowledgeUpdate) Description() string {
	return `
a knowledge base update interface.
use when you need to update knowledge base content.
your output should be in the following json format.
` + k.outPutParser.GetFormatInstructions()
}

func (k *KnowledgeUpdate) Call(ctx context.Context, input string) (string, error) {
	if err := k.svc.Auth(ctx); err != nil {
		return "", err
	}

	var data any
	data, err := k.outPutParser.Parse(input)
	if err != nil {
		// ```json str

		t := make(map[string]any)
		if err := json.Unmarshal([]byte(input), &t); err != nil {
			return "", err
		}

		return "", err
	}

	file := data.(map[string]any)

	f, err := os.Open(fmt.Sprintf("%v", file["path"]))
	if err != nil {
		return "", err
	}
	finfo, err := f.Stat()
	if err != nil {
		return "", err
	}

	p := documentloaders.NewPDF(f, finfo.Size())
	chunkedDocuments, err := p.LoadAndSplit(ctx, textsplitter.NewRecursiveCharacter(textsplitter.WithChunkSize(200),
		textsplitter.WithChunkOverlap(1)))
	if err != nil {
		return "", err
	}

	if k.store == nil {
		k.store, err = getKnowledgeStore(ctx, k.svc)
		if err != nil {
			return "", err
		}
	}

	_, err = k.store.AddDocuments(ctx, chunkedDocuments)
	if err != nil {
		return "", err
	}

	return Success, nil
}

func getKnowledgeStore(ctx context.Context, svc *svc.ServiceContext) (*redisvector.Store, error) {
	embedder, err := embeddings.NewEmbedder(svc.LLMs)
	if err != nil {
		return nil, err
	}

	return redisvector.New(ctx, redisvector.WithEmbedder(embedder), redisvector.WithConnectionURL("redis://"+svc.Config.
		Redis.
		Addr), redisvector.WithIndexName("knowledge", true))
}
